import os
import pandas as pd
from pathlib import Path
import dolphindb as ddb
import traceback
from datetime import datetime, timedelta
import efinance as ef

class RealtimeImporter:
    """实时期货数据导入器"""
    
    def __init__(self):
        self.session = None
        self.upserter = None
        
    def connect_db(self):
        """连接DolphinDB数据库"""
        self.session = ddb.session()
        self.session.connect("localhost", 8848, "admin", "123456")
        self.upserter = ddb.TableUpserter(
            dbPath="dfs://futures", 
            tableName="daily_all", 
            ddbSession=self.session,
            keyColNames=["date", "code", "exchange_code"]
        )

    def _get_exchange_info(self, quote_id):
        """根据行情ID获取交易所信息"""
        exchange_map = {
            '113': {'code': 'SHFE', 'name': '上期所'},
            '114': {'code': 'DCE', 'name': '大商所'},
            '115': {'code': 'ZCE', 'name': '郑商所'},
            '116': {'code': 'CFFEX', 'name': '中金所'},
            '142': {'code': 'INE', 'name': '上海能源期货交易所'}
        }
        
        prefix = quote_id.split('.')[0]
        return exchange_map.get(prefix, {'code': None, 'name': None})

    def get_all_futures_data(self):
        """获取所有期货合约的日线数据"""
        try:
            # 获取所有期货合约信息
            base_info = ef.futures.get_futures_base_info()
            
            # 根据行情ID前缀过滤出三个主要交易所的合约
            base_info['prefix'] = base_info['行情ID'].str.split('.').str[0]
            needed_exchanges = ['113', '114', '115']  # 上期所、大商所、郑商所
            base_info = base_info[base_info['prefix'].isin(needed_exchanges)]
            base_info = base_info.drop('prefix', axis=1)  # 删除临时列
            
            # 过滤掉主力连续合约和次主力连续合约
            base_info = base_info[~base_info['期货代码'].str.contains('m$|M$|s$|S$', regex=True)]
            
            # 获取当前日期
            today = datetime.now()
            start_date = (today - timedelta(days=365)).strftime('%Y%m%d')
            end_date = today.strftime('%Y%m%d')
            
            # 创建行情ID与交易所信息的映射
            quote_info = {}
            for _, row in base_info.iterrows():
                quote_id = row['行情ID']
                code = row['期货代码']
                exchange_info = self._get_exchange_info(quote_id)
                quote_info[code] = exchange_info
            
            # 获取所有合约的行情数据
            quote_ids = base_info['行情ID'].tolist()
            
            df_dict = ef.futures.get_quote_history(
                quote_ids,
                beg=start_date,
                end=end_date,
                klt=101,  # 日线数据
                return_df=True  # 返回DataFrame而不是dict
            )
            
            if df_dict is None or df_dict.empty:
                return None
            
            # 添加交易所信息
            df_dict['exchange_code'] = df_dict['期货代码'].map(lambda x: quote_info.get(x, {}).get('code'))
            df_dict['exchange_name'] = df_dict['期货代码'].map(lambda x: quote_info.get(x, {}).get('name'))
            
            # 重命名列以匹配数据库结构
            rename_dict = {
                '日期': 'date',
                '期货代码': 'code',
                '期货名称': 'name',
                '开盘': 'open',
                '最高': 'high',
                '最低': 'low',
                '收盘': 'close',
                '成交量': 'volume',
                '成交额': 'amount'
            }
            df_dict = df_dict.rename(columns=rename_dict)
            
            # 转换日期格式
            df_dict['date'] = pd.to_datetime(df_dict['date'])
            
            # 按合约代码分组，计算前一天的收盘价作为pre_close
            df_dict['pre_close'] = df_dict.groupby('code')['close'].shift(1)
            
            # 对于第一个交易日，使用当天的收盘价作为pre_close
            df_dict['pre_close'] = df_dict['pre_close'].fillna(df_dict['close'])
            
            # 选择需要的列
            final_columns = ['date', 'code', 'name', 'exchange_code', 'exchange_name',
                            'high', 'low', 'open', 'close', 'pre_close', 'volume', 'amount']
            
            result_df = df_dict[final_columns]
            
            return result_df
            
        except Exception as e:
            print(f"获取数据失败: {str(e)}")
            print(f"详细错误信息: {traceback.format_exc()}")
            return None

    def import_to_db(self, df):
        """导入数据到数据库"""
        try:
            self.connect_db()
            self.upserter.upsert(df)
            print(f"成功导入 {len(df)} 条数据")
        except Exception as e:
            print(f"导入数据失败: {str(e)}")
        finally:
            if self.session:
                self.session.close()

    def update_data(self):
        """更新数据"""
        print(f"\n开始更新期货数据... {datetime.now()}")
        df = self.get_all_futures_data()
        if df is not None and not df.empty:
            self.import_to_db(df)
        print(f"数据更新完成 {datetime.now()}")

def main():
    """手动更新数据"""
    importer = RealtimeImporter()
    importer.update_data()

if __name__ == "__main__":
    main() 